"""
  Copyright (c) 2024 Northeastern University (China)
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
"""
import time

import pymongo
from pymongo import errors

import logging
import os
from urllib.parse import quote

import sys

from gridfs import GridFS
from tqdm import tqdm

ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(ROOT_DIR)
from src.common_utils.configure import CONFIG

dblogger = logging.getLogger('mongodb')
dblogger.setLevel(logging.INFO)
log_file = os.path.join(os.getcwd(), 'mongodb_log.txt')
file_handler = logging.FileHandler(log_file, mode='w+')
file_handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s  %(filename)s : %(levelname)s  %(message)s', datefmt='%a, %d %b %Y %H:%M:%S')
file_handler.setFormatter(formatter)
dblogger.addHandler(file_handler)


class MongoDB:
    last_position = 0#游标位置，特地为yield返回设置的游标位置
    def __init__(self, db_name, collection_name):
        self.client = pymongo.MongoClient(self.get_mongourl())
        self.db = self.client[db_name]
        self.collection = self.db[collection_name]
        self.session = self.client.start_session()
        # 初始化GridFS对象
        self.gridfs = GridFS(self.db, collection=collection_name)
        # if collection_name=="OSSDB":
        #     self.collection.create_index([("oss_name", pymongo.DESCENDING), ("oss_source", pymongo.ASCENDING)],
        #                                  unique=True)
        if collection_name == "InitialSig" or collection_name == "Function_Src" or  collection_name == "TagInfo"\
                or  collection_name=="Component" or collection_name=="VerWeight":
            index_name = "sig_full_name_ - 1"  # 索引的名称
            index_info = self.collection.index_information()
            if index_name not in index_info:
                self.collection.create_index([("sig_full_name", pymongo.DESCENDING)], unique=True)
        if collection_name == "RepoFunction":
            index_name = "sig_full_tag_name_ - 1"  # 索引的名称
            index_info = self.collection.index_information()
            if index_name not in index_info:
                self.collection.create_index([("sig_full_tag_name", pymongo.DESCENDING)], unique=True)
        if collection_name == "FuncMap":
            index_name = "func_hash_ - 1"  # 索引的名称
            index_info = self.collection.index_information()
            if index_name not in index_info:
                self.collection.create_index([("func_hash", pymongo.DESCENDING)], unique=True)
        if collection_name == "InitialSig_beta" or collection_name == "Function_Src_beta" or collection_name == "TagInfo_beta" \
                or collection_name == "Component_beta" or collection_name=="VerWeight":
            index_name = "sig_full_name_ - 1"  # 索引的名称
            index_info = self.collection.index_information()
            if index_name not in index_info:
                self.collection.create_index([("sig_full_name", pymongo.DESCENDING)], unique=True)
        ###########testdata##########
        if collection_name == "harmonyProjects_build_srcfile":
            index_name = "filepath_ - 1"  # 索引的名称
            index_info = self.collection.index_information()
            if index_name not in index_info:
                self.collection.create_index([("filepath", pymongo.DESCENDING)], unique=True)
        if collection_name=='Indirect':
            A = "A_ - 1"  # 索引的名称
            B ="B_ -1"
            index_info = self.collection.index_information()
            if A not in index_info and B not  in index_info:
                self.collection.create_index([("A", pymongo.DESCENDING),("B", pymongo.DESCENDING)], unique=True)




    def get_mongourl(self) -> str:
        ip = CONFIG['mongodb']['ip']
        port = CONFIG['mongodb']['port']
        user = quote(CONFIG['mongodb']['user'])
        pwd = quote(CONFIG['mongodb']['pwd'])
        if user:
            return (f'mongodb://{user}:{pwd}@{ip}:{port}/?authSource=HarmonyOSS')
        else:
            return f'mongodb://{ip}:{port}/'

    # def get_large_data(self,filterDict={}, fieldDict={}, count: int = 0,db_name="",collection_name=""):
    #         cursor = self.collection.find(filterDict,fieldDict,no_cursor_timeout=True,session=self.session).batch_size(30).sort("$natural", 1).skip(count)
    #         for row in cursor:
    #            yield row
    def get_large_data(self, filterDict={}, fieldDict={}, count=0,db_name="",collection_name=""):
        self.last_position=0
        document_count = self.get_document_count(filterDict)
        while self.last_position+count< document_count:
            try:
                cursor = self.collection.find(filterDict, fieldDict, no_cursor_timeout=True,
                                              session=self.session).batch_size(30).sort("$natural", 1).skip( count + self.last_position)
                for row in cursor:
                    self.last_position += 1
                    yield row
            except (pymongo.errors.CursorNotFound, pymongo.errors.AutoReconnect) as e:
                print(f"游标超时错误: {e}")
                print("等待5秒重新连接数据库...")
                self.client.close()
                time.sleep(5)  # 等待一段时间，避免频繁重连
                self.client = pymongo.MongoClient(self.get_mongourl())
                self.db = self.client[db_name]
                self.collection = self.db[collection_name]
                self.session = self.client.start_session()
                # 初始化GridFS对象
                self.gridfs = GridFS(self.db, collection=collection_name)
                continue
    # def get_large_data(self, filterDict={}, fieldDict={}, count=0,db_name="",collection_name=""):
    #     cursor = self.collection.find(filterDict, fieldDict, no_cursor_timeout=True,
    #                                   session=self.session).batch_size(30).sort("$natural", 1).skip(count+self.last_position)
    #     document_count = self.get_document_count(filterDict)
    #     while self.last_position+count < document_count:
    #         try:
    #             for row in cursor:
    #                 self.last_position += 1
    #                 yield row
    #         except (pymongo.errors.CursorNotFound, pymongo.errors.AutoReconnect) as e:
    #             print(f"遇到CursorNotFound或AutoReconnect错误: {e}")
    #             print("正在重新查询...")
    #             # 在异常处理代码中重新创建游标
    #             cursor.close()
    #             cursor = self.collection.find(filterDict, fieldDict, no_cursor_timeout=True,session=self.session).batch_size(30).sort("$natural", 1).skip(count+self.last_position)
    def get_document_count(self,filterDict={}):
        # 获取集合中的数据数量
        count = self.collection.count_documents(filterDict)
        return count


    def add_many(self, adv_info_list):
        result = None
        try:
            result = self.collection.insert_many(adv_info_list)
            dblogger.info("add many:{}!".format(len(result.inserted_ids)))
        except Exception as e:
            dblogger.error(f"add many failed :{str(e)}!")
        return result
    def update_many(self,operation_list):
        result=None
        try:
            # 执行批量更新操作
            result = self.collection.bulk_write(operation_list)
            dblogger.info(f"Updated many documents: {result.modified_count} .")
        except Exception as e:
            dblogger.error(f"Updated many failed :{str(e)}!")
        return result
    def clear_collection(self):
        self.collection.delete_many({})

    def add_many_update_repeat(self, adv_info_list):
        for doc in adv_info_list:
            try:
                result=self.collection.insert_one(doc)
                result = result.inserted_id
                dblogger.info("add one:{}!".format(result))
            except pymongo.errors.DuplicateKeyError:
                # 主键冲突，执行更新操作
                query = {'_id': doc['_id']}  # 使用正确的主键查询条件
                update = {'$set': doc}  # 更新的操作，这里将整个文档替换为新值
                try:
                    self.collection.update_one(query, update)
                except Exception as e:
                    dblogger.error(f"update one failed :{str(e)}!")

    def add_one(self, adv_info):
        result = None
        try:
            result = self.collection.insert_one(adv_info)
            result = result.inserted_id
            dblogger.info("add one:{}!".format(result))
        except Exception as e:
            dblogger.error(f"add one failed :{str(e)}!")
        return result
    def delete_one(self,condition):
        try:
            result=self.collection.delete_one(condition)
            dblogger.info("delete one:{}!".format(result.deleted_count))
        except Exception as e :
            dblogger.error(f"delete one failed :{str(e)}!")
    def add_one_update_repeat(self,adv_info):
        result = None
        try:
            result = self.collection.insert_one(adv_info)
            result = result.inserted_id
            dblogger.info("add one:{}!".format(result))
        except pymongo.errors.DuplicateKeyError:
            # 主键冲突，执行更新操作
            query = {'_id': adv_info['_id']}  # 使用正确的主键查询条件
            update = {'$set': adv_info}  # 更新的操作，这里将整个文档替换为新值
            try:
                self.collection.update_one(query, update)
            except Exception as e:
                dblogger.error(f"update one failed :{str(e)}!")
        return result


    def get(self, filterDict={}, fieldDict={},count:int=0):
        ret = []
        cur = self.collection.find(filterDict, fieldDict).sort("$natural", 1).skip(count)
        dblogger.info("get data:{}!".format(cur))
        for item in cur:
            ret.append(item)
        return ret

    def update_one_with_id(self, _id, content={}):
        try:
            result = self.collection.update_one(filter={"_id": _id}, update={'$set': content})
            dblogger.info("update data:{}!".format(result))
            return result
        except Exception as e:
            dblogger.error(f"update one failed :{str(e)}!")
            return None

    def update_one_with_id_push(self, _id, content={}):
        result = self.collection.update_one(filter={"_id": _id}, update={'$push': content})
        return result

    def update_one_with_condition(self, condition, content={}):
        try:
            result = self.collection.update_one(filter=condition, update={'$set': content})
            dblogger.info("update one one:{}!".format(result))
            return result
        except Exception as e:
            dblogger.error(f"update one failed :{str(e)}!")
            return None
    def use_Gridfs_store(self,file,filename):
        try:
            file_id=self.gridfs.put(file,filename=filename)
            dblogger.info("use_Gridfs_store one:{}!".format(file_id))
            return file_id
        except Exception as e:
            dblogger.error(f"use_Gridfs_store failed :{str(e)}!")
            return None

    def use_Gridfs_delete(self, fileid):
        try:
            self.gridfs.delete(fileid)
            dblogger.info(f"use_Gridfs_delete one:{fileid}!")
        except Exception as e:
            dblogger.error(f"use_Gridfs_delete failed :{str(e)}!")
            return None
    def use_Gridfs_get(self,file_id):
        try:
            # 从GridFS中检索文件
            file = self.gridfs.get(file_id)
            # 检查文件是否存在
            if file is not None:
                # 读取文件内容
                data = file.read()
                dblogger.info("use_Gridfs_get one:{}!".format(file_id))
                return data
            else:
                dblogger.info("use_Gridfs_get one:None!")
                return None
        except Exception as e:
            dblogger.error(f"use_Gridfs_get failed :{str(e)}!")
            return None
